"""Adapter 基类定义"""

from __future__ import annotations

from abc import ABC, abstractmethod
from typing import Any, Dict, List

import numpy as np


class AdapterBackend(ABC):
    """
    推理适配器抽象基类
    
    所有后端适配器（ONNX Runtime、TensorRT、SophonTPU 等）必须继承此类
    并实现所有抽象方法，以确保接口一致性。
    
    符合规范：docs/adapter-spec.md
    """
    
    @staticmethod
    @abstractmethod
    def is_available() -> bool:
        """
        检查适配器是否可用
        
        Returns:
            bool: 适配器是否可用（依赖已安装、硬件可用等）
        
        实现要点：
            1. 检查必要的依赖包是否已安装
            2. 检查硬件设备是否可用（如适用）
            3. 应该是轻量级检查，不应进行耗时的初始化操作
            4. 如果所有依赖都满足，返回 True，否则返回 False
        
        示例:
            @staticmethod
            def is_available() -> bool:
                try:
                    import onnxruntime
                    return True
                except ImportError:
                    return False
        """
        pass
    
    @abstractmethod
    def prepare(self, model_spec: Dict[str, Any]) -> Any:
        """
        加载模型并准备推理环境
        
        Args:
            model_spec: 模型规格字典，包含以下字段：
                - source_onnx (str): 官方 ONNX 模型文件路径（绝对路径）
                - name (str): 模型名称
                - sha256 (str): 模型文件 SHA256 哈希值（用于合规校验）
                - inputs (Dict): 输入规格 {input_name: {shape, dtype, layout, range}}
                - outputs (List[str]): 输出节点名称列表
                - task (str): 任务类型（classification/segmentation/nlp）
                - opset (int): ONNX opset 版本
        
        Returns:
            推理句柄（handle），类型由具体实现决定
            - ONNX Runtime: onnxruntime.InferenceSession
            - TensorRT: tensorrt.Engine
            - SophonTPU: bmrt.Runtime
            等等
        
        Raises:
            ModelLoadError: 模型加载失败
            ComplianceError: 模型哈希校验失败
        
        实现要点：
            1. 必须使用 model_spec["source_onnx"] 指定的官方 ONNX 模型
            2. 可选：在 prepare() 阶段转换为后端专用格式（如 bmodel、om、trt 等）
            3. 转换必须保证模型拓扑、算子类型、IO 规格与原始 ONNX 一致
            4. 建议：校验模型文件 SHA256，不匹配则抛出 ComplianceError
            5. 可进行后端特定优化（如图优化、算子融合），但不得改变模型语义
        """
        pass
    
    @abstractmethod
    def infer(self, handle: Any, inputs: Dict[str, np.ndarray]) -> List[np.ndarray]:
        """
        执行单次推理
        
        Args:
            handle: prepare() 返回的推理句柄
            inputs: 输入张量字典 {input_name: np.ndarray}
                - 每个 np.ndarray 的 shape、dtype 应符合模型规格
                - batch 维度必须一致
        
        Returns:
            输出张量列表 List[np.ndarray]
            - 按模型输出节点顺序返回
            - 每个输出都是 numpy array（即使后端使用其他格式，也需转换回 numpy）
        
        Raises:
            InferenceError: 推理执行失败
        
        实现要点：
            1. 接受 numpy 格式输入，返回 numpy 格式输出
            2. 内部可转换为后端特定格式（如 CUDA Tensor、TPU Tensor）
            3. 必须保证数值精度满足合规要求（见 docs/requirements.md）
        """
        pass
    
    @abstractmethod
    def teardown(self, handle: Any) -> None:
        """
        释放推理句柄和相关资源
        
        Args:
            handle: prepare() 返回的推理句柄
        
        实现要点：
            1. 释放模型占用的内存、显存等资源
            2. 关闭设备连接（如适用）
            3. 某些后端（如 ONNX Runtime）可能由 GC 自动管理，此时可以为空实现
        """
        pass
    
    @abstractmethod
    def get_metadata(self) -> Dict[str, Any]:
        """
        返回适配器和后端的元信息
        
        Returns:
            元信息字典，必须包含以下字段：
                - adapter_backend (str): 适配器名称（如 "onnxruntime", "tensorrt", "sophon_tpu"）
                - adapter_version (str): 适配器版本
                - backend (str): 后端名称
                - backend_version (str): 后端版本
                - device (str): 设备类型（如 "CUDA", "CPU", "TPU", "NPU"）
                - providers (List[str]): 执行提供者列表（如适用）
            
            可选字段：
                - cuda_version (int): CUDA 版本（如使用 CUDA）
                - conversion_tool (str): 模型转换工具（如转换为 bmodel、om 等）
                - optimization_flags (List[str]): 优化选项
                - bmodel_path (str): 转换后的模型路径（如适用）
                等等
        
        实现要点：
            1. 此方法可能在 prepare() 之前或之后调用
            2. 如果 prepare() 未调用，应返回后端的通用信息
            3. 如果 prepare() 已调用，应返回实际使用的配置信息
        """
        pass
    
    @staticmethod
    def compute_sha256(file_path: str) -> str:
        """
        计算文件的 SHA256 哈希值（工具方法）
        
        Args:
            file_path: 文件路径
        
        Returns:
            SHA256 哈希值（十六进制字符串）
        """
        import hashlib
        
        sha256_hash = hashlib.sha256()
        with open(file_path, "rb") as f:
            # 分块读取，避免大文件内存溢出
            for byte_block in iter(lambda: f.read(4096), b""):
                sha256_hash.update(byte_block)
        return sha256_hash.hexdigest()
    
    def validate_model_compliance(self, model_spec: Dict[str, Any]) -> None:
        """
        验证模型合规性（工具方法，子类可选实现）
        
        Args:
            model_spec: 模型规格字典
        
        Raises:
            ComplianceError: 模型不合规
        
        默认实现：仅校验 SHA256
        子类可覆盖此方法以添加更多校验（如算子子集检查、拓扑验证等）
        """
        from ...exceptions import ComplianceError
        
        model_path = model_spec.get("source_onnx")
        expected_sha256 = model_spec.get("sha256")
        
        if expected_sha256 and model_path:
            actual_sha256 = self.compute_sha256(model_path)
            if actual_sha256 != expected_sha256:
                raise ComplianceError(
                    f"Model hash mismatch for {model_path}\n"
                    f"Expected: {expected_sha256}\n"
                    f"Actual: {actual_sha256}"
                )

